import numpy as np
from bishe_benchmark import get_all_nguyen_equations, generate_nguyen_data
from bishe_benchmark_use import use_GPlearn, use_NeuralGPlearn, use_NeuralPySR, use_PySR
import torch
import time
from typing import Dict
import json
import os, pandas as pd

def run_benchmark():
    """运行基准测试"""
    # 获取所有Nguyen方程
    equations = get_all_nguyen_equations()
 
    for func, name in equations:
        print(f"\n开始测试方程: {name}")

        # 生成数据
        df = pd.read_csv(f"nguyen_data/{name}_data.csv")
        X = df['X'].values
        y = df['y'].values
        X = X.reshape(-1, 1)
        y = y.reshape(-1, 1)
        # 优化超参数
        print(f"$$$$name,mse,train_time,func\n")
        # 使用PySR
        print("$$$$start PySR")
        _func, mse, train_time = use_PySR(X, y)
        # print(func, mse, train_time)
        print(f"$$$${name},{mse},{train_time},{_func}\n")
    


if __name__ == "__main__":
    """主函数"""
    # 设置随机种子
    np.random.seed(42)
    torch.manual_seed(42)
    run_benchmark()